#!/usr/bin/python
# -*- coding: utf-8 -*-
import socket
import sys
import threading
import Queue
import cPickle

# Base communication channel class. This class
# manages the comunication among the peers. 

# This server is multi-threaded, and it finishes its
# execution when the peers close all receive channels
# and it closes all send channels (with an
# explicit server.close_channel()

class origin_msg(object):
    def __init__(self, pid):
        self.pid = pid

class comm_channel(object):
    def __init__(self, pid, hosts, dispatcher):
        self.pid = pid
        self.hosts = hosts
        self.dispatcher = dispatcher
        self.channels = {}

        self.SEPARATOR = "XdemX"

        self.partners = [x for x in hosts.keys() if x != self.pid]

        for i in self.partners:
            self.channels[i] = True
    
    def print_output(self, text):
        pass

    def close_channel(self, dest):
        self.send_message(dest, None) 

    def channel_state(self, channel):
        if channel in self.channels:
            return self.channels[channel]
        else:
            return False

    def get_hosts(self):
        return self.hosts

    def get_pid(self):
        return self.pid
    
    def send_message(self, dest, msg):
        pass
        
    def recv_message(self):
        return None

    def network_size(self):
        return len(self.hosts)

    def close(self):
        pass

class TCP_comm_channel(comm_channel):

    MSG_SIZE_LIMIT_ = 4096
    def __init__(self, pid, hosts, dispatcher, 
                 socket=socket, threading=threading):
        comm_channel.__init__(self, pid, hosts, dispatcher)

        # Queue variables
        self.socket_ = socket
        self.in_queue = Queue.Queue(0)
        self.out_queue = {}
        self.connected = {}

        self.sck = self.socket_.socket (self.socket_.AF_INET, 
                                        self.socket_.SOCK_STREAM)
        self.sck.bind(self.hosts[self.pid])
        self.sck.listen(10)

        self.dc = self.socket_.socket(self.socket_.AF_INET, 
                                      self.socket_.SOCK_STREAM)
        self.dc.connect(self.dispatcher)
        self.dc.recv(self.MSG_SIZE_LIMIT_)
        
        for i in self.partners:
            self.out_queue[i] = Queue.Queue(0)
            self.connected[i] = False

        for i in self.partners:
            th_send = threading.Thread(
                target=self.send_channel,
                name='PID: %d - Send %d' % (self.pid, i),
                args=[i])
            th_send.start()

        for i in self.partners:
            (conn, addr) = self.sck.accept()
            th_recv = threading.Thread(
                target=self.recv_channel,
                name='PID: %d - Receive %s' % (self.pid, addr),
                args=(conn, addr))
            th_recv.start()

    def send_message(self, dest, msg):
        self.out_queue[dest].put(msg)

    def recv_message(self):
        origin, msg = self.in_queue.get()
        return (origin, msg)

    def send_channel(self, *args):
        try:
            dest = args[0]
            msg = None
            s = self.socket_.socket(self.socket_.AF_INET, 
                                    self.socket_.SOCK_STREAM)
            s.connect(self.hosts[dest])
            obj = origin_msg (self.pid)
            msg = ''.join([cPickle.dumps(obj), self.SEPARATOR])
            size = s.send(msg)
            while (size != len(msg)):
                msg = msg[size:]
                size = s.send(msg)
                
            while(obj != None):
                obj = self.out_queue[dest].get()
                if obj != None:
                    full_msg = ''.join([cPickle.dumps(obj), 
                                        self.SEPARATOR])
                    for i in xrange(0, len(full_msg) / 
                                    self.MSG_SIZE_LIMIT_ + 1):
                        msg = full_msg[i * self.MSG_SIZE_LIMIT_: 
                                       (i + 1) * self.MSG_SIZE_LIMIT_]
                        size = s.send(msg)
                        while (size != len(msg)):
                            msg = msg[size:]
                            size = s.send(msg)
        finally:
            s.close()
        
    def recv_channel(self, *args):
        try:
            conn, addr = args
            origin = -1
            history = []
            new_data = conn.recv(self.MSG_SIZE_LIMIT_)
            history.append(new_data)
            while (new_data and (len(history) != 0)):
                if new_data.find(self.SEPARATOR) != -1:
                    data = ''.join(history)
                    pickled_objs = data.split(self.SEPARATOR)
                    history = [pickled_objs[-1]]
                    pickled_objs = pickled_objs[:-1]
                    for i in pickled_objs:
                        i.replace("\\n","\n") 
                        msg = cPickle.loads(i)
                        if isinstance(msg, origin_msg):
                            origin = msg.pid
                            self.connected[origin] = True
                        elif origin != -1:
                            self.in_queue.put((origin, msg))
                new_data = conn.recv(self.MSG_SIZE_LIMIT_)
                if new_data:
                    history.append(new_data)
        finally:
            conn.close()

    def print_output(self, text):
        text = "Server %d: %s \n" % (self.pid, text)
        size = self.dc.send(text)
        while (size != len(text)):
            text = text[size:]
            size = self.dc.send(text)

    def close(self):
        self.dc.close()
